package ent

import (
	"context"
	"time"

	"google.golang.org/grpc/grpclog"

	"entgo.io/ent/dialect"
)

var _ dialect.Driver = (*MyDriver)(nil)

func NewMyDriver(dri dialect.Driver, log grpclog.LoggerV2) dialect.Driver {
	return &MyDriver{driver: dri, log: log}
}

type MyDriver struct {
	driver dialect.Driver
	log    grpclog.LoggerV2
}

func (m *MyDriver) Exec(ctx context.Context, query string, args, v interface{}) error {
	s := time.Now()
	err := m.driver.Exec(ctx, query, args, v)
	m.log.Infof("[%.3fms] %s", float64(time.Since(s).Nanoseconds())/1e6, ExplainSQL(query, nil, `"`, args))
	return err
}

func (m *MyDriver) Query(ctx context.Context, query string, args, v interface{}) error {
	s := time.Now()
	err := m.driver.Query(ctx, query, args, v)
	m.log.Infof("[%.3fms] %s", float64(time.Since(s).Nanoseconds())/1e6, ExplainSQL(query, nil, `"`, args))
	return err
}

func (m *MyDriver) Tx(ctx context.Context) (dialect.Tx, error) {
	tx, err := m.driver.Tx(ctx)
	if err != nil {
		return nil, err
	}
	return &MyTx{tx: tx, log: m.log}, nil
}

func (m *MyDriver) Close() error {
	return m.driver.Close()
}

func (m *MyDriver) Dialect() string {
	return m.driver.Dialect()
}

var _ dialect.Tx = (*MyTx)(nil)

type MyTx struct {
	tx  dialect.Tx
	log grpclog.LoggerV2
}

func (m *MyTx) Exec(ctx context.Context, query string, args, v interface{}) error {
	s := time.Now()
	err := m.tx.Exec(ctx, query, args, v)
	m.log.Infof("[%.3fms] %s", float64(time.Since(s).Nanoseconds())/1e6, ExplainSQL(query, nil, `"`, args))
	return err
}

func (m *MyTx) Query(ctx context.Context, query string, args, v interface{}) error {
	s := time.Now()
	err := m.tx.Query(ctx, query, args, v)
	m.log.Infof("[%.3fms] %s", float64(time.Since(s).Nanoseconds())/1e6, ExplainSQL(query, nil, `"`, args))
	return err
}

func (m *MyTx) Commit() error {
	return m.tx.Commit()
}

func (m *MyTx) Rollback() error {
	return m.tx.Rollback()
}
